#!/bin/bash

# Initialize optional parameters
TCPD=""
OUTPUT=""
DAEMON=""

# Parse command-line options
while getopts "f:m:p:w:t:o:-:" opt; do
  case $opt in
    f) MODEL_SCRIPT="$OPTARG"
    ;;
    m) MODEL_FILE="$OPTARG"
    ;;
    p) MODEL_PORT="$OPTARG"
    ;;
    w) MAPPED_PORT="$OPTARG"
    ;;
    t) TCPD="$OPTARG"
    ;;
    o) OUTPUT="$OPTARG"
    ;;
    -)
      case "${OPTARG}" in
        daemo) DAEMON="--daemo"
        ;;
        *) echo "Invalid option: --${OPTARG}" >&2
        ;;
      esac
    ;;
    *) echo "Invalid option: -$OPTARG" >&2
    ;;
  esac
done

if [ -z "$MODEL_FILE" ] || [ -z "$MODEL_PORT" ] || [ -z "$MAPPED_PORT" ]; then
  echo "Usage: $0 -f model.py -m reconet.ckpt -p 5610 -w tcp://localhost:5553 [-t tcpd] [-o output] [--daemo]"
  exit 1
fi

MAPPED_PORT_NUM=$(echo $MAPPED_PORT | sed 's/.*://')

# Function to start the model script
start_model_script() {
  if [[ $MODEL_FILE == *.om ]]; then
    echo "Detected .om model file. Starting airloader..."
    airloader -m $MODEL_FILE -p $MODEL_PORT &
  else
    if [ -z "$MODEL_SCRIPT" ]; then
      echo "Model script (-f) is required for non-.om model files."
      exit 1
    fi
    echo "Starting model script..."
    python3.9 -c "
import sys
import importlib

model_file = '$MODEL_SCRIPT'
ckpt_file = '$MODEL_FILE'
port = '$MODEL_PORT'

if '-p' not in sys.argv and '--port' not in sys.argv:
    sys.argv.append('-p')
    sys.argv.append(port)
from mindspore_model import MindsporeModel   
module_name = model_file.replace('.py', '')
model_module = importlib.import_module(module_name)
net = model_module.ReCoNet()
MindsporeModel(net, ckpt_file).run()
" &
  fi

  # Store the PID of the model script
  MODEL_PID=$!
}

# Function to start the RECOapplet script
start_reco_applet() {
  echo "Starting RECOapplet..."
  python3.9 RECOapplet.py -r tcp://localhost:$MODEL_PORT -p $MAPPED_PORT_NUM ${TCPD:+-t $TCPD} ${OUTPUT:+-o $OUTPUT} $DAEMON &

  # Store the PID of the RECOapplet script
  APPLET_PID=$!
}

# Function to kill both processes
cleanup() {
  echo "Stopping processes..."
  kill $MODEL_PID
  kill $APPLET_PID
}

# Trap SIGINT and SIGTERM to cleanup
trap cleanup SIGINT SIGTERM

# Start the model script in the background
start_model_script

# Start the RECOapplet script in the background
start_reco_applet

# Wait for the model script to start
wait $MODEL_PID

# Wait for both processes to finish
wait $APPLET_PID
